import numpy as np
import matplotlib.pyplot as plt
import numpy.random as npr
import matplotlib
matplotlib.use("TkAgg")

def experiment():
    r, p = 5, 0.3    # 需要的成功次数 = r, 每次成功概率 = p
    trials = 100_000

    # 模拟数据：X = 达到 r 次成功前的失败次数
    samples_neg_binom = npr.negative_binomial(n=r, p=p, size=trials)

    # 理论均值
    theoric_mean = r * (1-p) / p
    empirical_mean = samples_neg_binom.mean()

    # 画图
    plt.figure(figsize=(6,4))
    max_val = np.max(samples_neg_binom)
    plt.hist(samples_neg_binom, bins=range(0, max_val+2), density=True,
             alpha=0.6, color='skyblue', edgecolor='black')
    plt.axvline(theoric_mean, color='red', linestyle='--',
                label=f'Theoretical E[X]={theoric_mean:.2f}')
    plt.axvline(empirical_mean, color='green', linestyle=':',
                label=f'Empirical mean={empirical_mean:.2f}')
    plt.title(f'Negative Binomial Distribution (r={r}, p={p})')
    plt.xlabel("Number of failures before r successes")
    plt.ylabel("Probability")
    plt.legend()
    plt.tight_layout()
    plt.show()

# 测试 1：直观例子
def test01():
    r, p = 3, 0.5
    trials = 20
    samples = npr.negative_binomial(n=r, p=p, size=trials)
    print(samples)

# 测试 2：现实场景
'''
例子：机器需要 3 次成功启动 (r=3)，
每次启动成功概率 0.2。
问：达到 3 次成功前恰好 5 次失败的概率是多少？
'''
def test02():
    r, p = 3, 0.2
    # 理论概率
    from math import comb
    k = 5  # 恰好失败 5 次
    prob_theory = comb(r+k-1, k) * (1-p)**k * p**r
    # 模拟估计
    trials = 200000
    prob_emp = np.mean(npr.negative_binomial(r, p, size=trials) == k)
    print(f"Theoretical P(X={k})={prob_theory:.5f}, Empirical≈{prob_emp:.5f}")

if __name__ == '__main__':
    # experiment()
    # test01()
    test02()
